Claim callbacks through the Execution API before running them - #70927
Claim callbacks through the Execution API before running them#70927potiuk wants to merge 2 commits into
Conversation
|
@seanghaeli - weren't you working on this a while ago, or was that a different workload token? Found it. I think these two PRs are going to collide, or at least overlap: #71192 |
| }, | ||
| ) | ||
|
|
||
| callback.state = CallbackState.RUNNING |
There was a problem hiding this comment.
Doesn't this conflict with the executors which already change the state from QUEUED to RUNNING?
LocalExecutor queues the RUNNING event before calling run_workload(), so it would be RUNNING before it gets here and this if not QUEUED block would throw an exception, failing the callback by default before it ever gets run. The AWS executors all set the state to RUNNING in self.running_state() as well. I'm pretty sure Celery is the exception here and every other executor would start failing.
I suspect the executors need to be updated to not also change the state, or this if queued check needs to be expanded to if not terminal state, maybe?
|
Yes, I think this PR fully overlaps with #71192. Since Ash has already taken a look at it, let's land that one? In the linked PR, the token is more narrowly scoped to just the A gap in mine though is that it doesn't include API versioning bundle, so I'll pull that over. |
|
Added the versioning bundle code to #71192 |
A CeleryExecutor worker runs an ExecuteCallback workload by importing the
callback path and calling it in a supervised subprocess. Unlike a task, a
callback made no authenticated call to the Execution API before doing so, so the
workload token minted for it in ExecuteCallback.make was never redeemed anywhere
and the callback executed whether or not the token was valid.
Add POST /execution/callbacks/{callback_id}/run, scoped to token:workload +
cb:self, which atomically transitions the callback QUEUED -> RUNNING, and call it
from supervise_callback before the subprocess starts. The server validates the
token (workload scope, subject == this callback id), so a missing, forged, or
mismatched token is rejected before the callback body runs. The transition is
single-shot, so a redelivered or replayed message that reaches an already-claimed
callback is refused rather than run twice, and it gives callbacks the RUNNING
state the workload docstring already describes.
cb:self mirrors ct:self for connection tests; the supervisor ordering mirrors
supervise_connection_test. The new endpoint is versioned as 2026-08-01 so older
clients continue to negotiate an API version in which it does not exist.
f97f111 to
c0be079
Compare
|
Agreed — these are the same change, and since you've already pulled the versioning bundle across, #71192 is the one to land. Your token model is the better of the two: mine only forces the workload token to be validated once, whereas yours makes it single-use and narrows what the callback can do after the exchange. Closing this one as superseded. Two things worth carrying over before it goes:
Happy to port both over to your branch myself if that's easier — just say the word and I'll open a PR against it. One thing that came over from my branch and needs a fresh look on rebase: the bundle now has No need to port Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting |
|
Just rebased onto main — the conflicts were in the Execution API version registry and the generated Task SDK datamodels, both resolved (the generated file was re-generated rather than hand-edited, and comes out byte-identical). @amoghrajesh you review most of the Execution API and Task SDK work; would you take a look? One thing worth an opinion on while reviewing: this adds Execution API version Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting |
…dress review guidance from apache#70927 - Register AddCallbackRunEndpoint under the existing unreleased 2026-10-30 version instead of a new 2026-08-01 version dated below it (contributing-docs/19: new migrations ride an existing unreleased version). - Rename the callback:self security scope to cb:self, matching ti:self/ct:self. - Stop persisting executor-reported RUNNING events for callbacks: executors emit RUNNING optimistically before the workload runs, racing the single-use token exchange that owns the QUEUED -> RUNNING transition. - Add supervisor-level tests pinning that the token exchange happens before the callback subprocess starts and a rejected exchange prevents it.
…dress review guidance from apache#70927 - Register AddCallbackRunEndpoint under the existing unreleased 2026-10-30 version instead of a new 2026-08-01 version dated below it (contributing-docs/19: new migrations ride an existing unreleased version). - Rename the callback:self security scope to cb:self, matching ti:self/ct:self. - Stop persisting executor-reported RUNNING events for callbacks: executors emit RUNNING optimistically before the workload runs, racing the single-use token exchange that owns the QUEUED -> RUNNING transition. - Add supervisor-level tests pinning that the token exchange happens before the callback subprocess starts and a rejected exchange prevents it.
…dress review guidance from apache#70927 - Register AddCallbackRunEndpoint under the existing unreleased 2026-10-30 version instead of a new 2026-08-01 version dated below it (contributing-docs/19: new migrations ride an existing unreleased version). - Rename the callback:self security scope to cb:self, matching ti:self/ct:self. - Stop persisting executor-reported RUNNING events for callbacks: executors emit RUNNING optimistically before the workload runs, racing the single-use token exchange that owns the QUEUED -> RUNNING transition. - Add supervisor-level tests pinning that the token exchange happens before the callback subprocess starts and a rejected exchange prevents it.
3 participants
A CeleryExecutor worker runs an
ExecuteCallbackworkload by importing thecallback path and calling it in a supervised subprocess. Unlike a task or a
connection test, a callback made no authenticated call to the Execution API
before doing so — so the workload token minted for it in
ExecuteCallback.makewas never redeemed anywhere, and the callback executed whether or not the token
was valid. This closes that gap and gives callbacks the
RUNNINGstate theworkload state machine already describes (
QUEUED → RUNNING → SUCCESS/FAILED).Change
POST /execution/callbacks/{callback_id}/run, scopedtoken:workload+cb:self, that atomically transitions the callbackQUEUED → RUNNING.supervise_callbackcalls it before starting the subprocess — mirroringsupervise_connection_test, which calls the server before touching credentials.cb:selfenforcement mirrors the existingct:selffor connection tests.2026-08-01, so a client negotiating an older API version behavesas though the endpoint does not exist.
The server validates the token (workload scope, subject == this callback id), so
a missing or mismatched token is refused before the callback body runs. The
transition is single-shot, so a redelivered/replayed message that reaches an
already-claimed callback is refused rather than run twice.
Tests
test_callbacks.py— claim transitions QUEUED→RUNNING; 404; 422; and 409 on acallback already RUNNING/terminal (replay protection)
TestCbSelfScopeEnforcement— a token minted for another callback is rejected (403)test_token_scope_boundaries— the route is registered workload-onlyTestSuperviseCallbackClaimsBeforeExecuting— the claim runs before thesubprocess, and a rejected claim prevents it from starting. Verified: removing
the claim call makes both fail.
@vincbeck @o-nikolas — flagging you as the security-team folks most across this
area. Please pull in @ferruzzi and anyone else who worked on the workload /
callback execution path if I have missed them.
Was generative AI tooling used to co-author this PR?
Generated-by: Claude Opus 5 (1M context) following the guidelines at
https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions